package com.zszq.ms.bigmodeltest.utils;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.CellType;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.ss.usermodel.WorkbookFactory;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * 解析文件工具类
 *
 * @author ms
 * @since 2023-07-07
 */
public class FileParser {

    /**
     * 解析文件并返回解析内容
     *
     * @param is 文件输入流
     * @param fileName 文件名
     * @return 文件内容哦个你
     * @throws Exception .
     */
    public static List<String> parseFile(InputStream is, String fileName) throws Exception {
        String extension = getFileExtension(fileName);
        switch (extension.toLowerCase()) {
            case "pdf":
                return parsePdf(is);
            case "doc":
            case "docx":
                return parseWord(is);
            case "xls":
            case "xlsx":
                return parseExcel(is);
            default:
                throw new Exception("不支持的格式");
        }
    }

    /**
     * PDF解析
     *
     * @param is 文件流
     * @return 文件解析内容
     * @throws Exception .
     */
    private static List<String> parsePdf(InputStream is) throws Exception {
        List<String> pages = new ArrayList<>();
        try (PDDocument document = PDDocument.load(is)) {
            PDFTextStripper stripper = new PDFTextStripper();
            for (int page = 1; page <= document.getNumberOfPages(); ++page) {
                stripper.setStartPage(page);
                stripper.setEndPage(page);
                pages.add(stripper.getText(document));
            }
        }
        return pages;
    }

    /**
     * word类型文件解析
     *
     * @param is 文件流
     * @return 解析内容
     * @throws Exception .
     */
    private static List<String> parseWord(InputStream is) throws Exception {
        List<String> pages = new ArrayList<>();
        try (XWPFDocument doc = new XWPFDocument(is)) {
            XWPFWordExtractor extractor = new XWPFWordExtractor(doc);
            String text = extractor.getText();
            // 简单按分页符分割
            String[] pageContents = text.split("\f");
            for (String content : pageContents) {
                pages.add(content.trim());
            }
        }
        return pages;
    }

    /**
     * excel类型文件解析
     *
     * @param is 文件流
     * @return 解析内容
     * @throws Exception .
     */
    public static List<String> parseExcel(InputStream is) throws Exception {
        List<String> sheets = new ArrayList<>();
        Workbook workbook = WorkbookFactory.create(is);
        for (int i = 0; i < workbook.getNumberOfSheets(); i++) {
            Sheet sheet = workbook.getSheetAt(i);
            StringBuilder sb = new StringBuilder();
            for (Row row : sheet) {
                for (Cell cell : row) {
                    sb.append(getCellValue(cell)).append("\t");
                }
                sb.append("\n");
            }
            sheets.add(sb.toString());
        }
        workbook.close();
        return sheets;
    }

    private static String getCellValue(Cell cell) {
        if (cell.getCellType().equals(CellType.STRING)) {
            return cell.getStringCellValue();
        } else if (cell.getCellType().equals(CellType.NUMERIC)) {
            return String.valueOf(cell.getNumericCellValue());
        } else if (cell.getCellType().equals(CellType.BOOLEAN)) {
            return String.valueOf(cell.getBooleanCellValue());
        } else {
            return "";
        }
    }

    private static String getFileExtension(String fileName) {
        return fileName.substring(fileName.lastIndexOf(".") + 1);
    }
}
